{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 循环神经网络的简洁实现\n",
    "import d2lzh as d2l\n",
    "import math\n",
    "from mxnet import autograd, gluon, init, nd\n",
    "from mxnet.gluon import loss as gloss, nn, rnn\n",
    "import time\n",
    "\n",
    "(corpus_indices, char_to_idx, idx_to_char, vocab_size) = d2l.load_data_jay_lyrics()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 定义模型\n",
    "# 下面构造一个含单隐藏层、隐藏单元个数为256的循环神经网络层 rnn_layer，并对权重做初始化\n",
    "num_hiddens = 256\n",
    "rnn_layer = rnn.RNN(num_hiddens)\n",
    "rnn_layer.initialize()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1, 2, 256)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 调用rnn_layer的成员函数begin_state返回初始化的隐藏状态列表\n",
    "batch_size = 2\n",
    "# 返回一个形状为（隐藏层个数，批量大小，隐藏单元个数）的元素\n",
    "state = rnn_layer.begin_state(batch_size=batch_size)\n",
    "state[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((35, 2, 256), 1, (1, 2, 256))"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "num_steps = 35\n",
    "X = nd.random.uniform(shape=(num_steps, batch_size, vocab_size))\n",
    "Y, state_new = rnn_layer(X, state)\n",
    "Y.shape, len(state_new), state_new[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 继承Block类来定义一个完整的循环神经网络\n",
    "class RNNModel(nn.Block):\n",
    "    def __init__(self, rnn_layer, vocab_size, **kwargs):\n",
    "        super(RNNModel, self).__init__(**kwargs)\n",
    "        self.rnn = rnn_layer\n",
    "        self.vocab_size = vocab_size\n",
    "        self.dense = nn.Dense(vocab_size)\n",
    "    \n",
    "    def forward(self, inputs, state):\n",
    "        # 将输入转置成(num_steps, batch_size)后获取one-hot向量表示\n",
    "        X = nd.one_hot(inputs.T, self.vocab_size)\n",
    "        Y, state = self.rnn(X, state)\n",
    "        # 全连接层会首先将Y的形状变成(num_steps*batch_size, num_hiddens), 它的输出\n",
    "        # 形状为(num_steps*batch_size, vocab_size)\n",
    "        output = self.dense(Y.reshape((-1, Y.shape[-1])))\n",
    "        return output, state\n",
    "    \n",
    "    def begin_state(self, *args, **kwargs):\n",
    "        return self.rnn.begin_state(*args, **kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 训练模型嗯\n",
    "def predict_rnn_gluon(prefix, num_chars, model, vocab_size, ctx, idx_to_char,\n",
    "                      char_to_idx):\n",
    "    # 使用model的成员函数来初始化隐藏状态\n",
    "    state = model.begin_state(batch_size=1, ctx=ctx)\n",
    "    output = [char_to_idx[prefix[0]]]\n",
    "    for t in range(num_chars + len(prefix) - 1):\n",
    "        X = nd.array([output[-1]], ctx=ctx).reshape((1,1))\n",
    "        (Y, state) = model(X, state) # 前向计算不需要传入模型参数\n",
    "        if t < len(prefix) - 1:\n",
    "            output.append(char_to_idx[prefix[t + 1]])\n",
    "        else:\n",
    "            output.append(int(Y.argmax(axis=1).asscalar()))\n",
    "    return ''.join([idx_to_char[i] for i in output])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ctx = d2l.try_gpu()\n",
    "model = RNNModel(rnn_layer, vocab_size)\n",
    "model.initialize(force_reinit=True, ctx=ctx)\n",
    "predict_rnn_gluon('分开', 10, model, vocab_size, ctx, idx_to_char, char_to_idx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_and_predict_rnn_gluon(model, num_hiddens, vocab_size, ctx,\n",
    "                                corpus_indices, idx_to_char, char_to_idx,\n",
    "                                num_epochs, num_steps, lr, clipping_theta,\n",
    "                                batch_size, pred_period, pred_len, prefixes):\n",
    "    loss = gloss.SoftmaxCrossEntropyLoss()\n",
    "    model.initialize(ctx=ctx, force_reinit=True, init=init.Normal(0.01))\n",
    "    trainer = gluon.Trainer(model.collect_params(), 'sgd',\n",
    "                            {'learning_rate': lr, 'momentum': 0, 'wd': 0})\n",
    "    for epoch in range(num_epochs):\n",
    "        l_sum, n, start = 0.0, 0, time.time()\n",
    "        data_iter = d2l.data_iter_consecutive(\n",
    "            corpus_indices, batch_size, num_steps, ctx)\n",
    "        state = model.begin_state(batch_size=batch_size, ctx=ctx)\n",
    "        for X, Y in data_iter:\n",
    "            for s in state:\n",
    "                s.detach()\n",
    "            with autograd.record():\n",
    "                (output, state) = model(X, state)\n",
    "                y = Y.T.reshape((-1,))\n",
    "                l = loss(output, y).mean()\n",
    "            l.backward()\n",
    "            # 梯度裁减\n",
    "            params = [p.data() for p in model.collect_params().values()]\n",
    "            d2l.grad_clipping(params, clipping_theta, ctx)\n",
    "            trainer.step(1)  # 因为已经误差取过均值，梯度不再做平均\n",
    "            l_sum += l.asscalar() * y.size\n",
    "            n += y.size\n",
    "        \n",
    "        if (epoch + 1) % pred_period == 0:\n",
    "            print('epoch %d, perplexity %f, time %.2f sec' % (\n",
    "                epoch + 1, math.exp(l_sum / n), time.time() - start))\n",
    "            for prefix in prefixes:\n",
    "                print(' -', predict_rnn_gluon(\n",
    "                    prefix, pred_len, model, vocab_size, ctx, idx_to_char, char_to_idx))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 50, perplexity 80.052697, time 0.25 sec\n",
      " - 分开 我不能这生 我不定这生 我不能这生 我不能这生 我不定这生 我不能这生 我不能这生 我不能这生 我\n",
      " - 不分开 我想我这生写谁一的我想就你我的见不就人 想要的让我疯狂的可爱女人 坏坏的让我疯狂的可爱女人 坏坏的\n",
      "epoch 100, perplexity 13.783161, time 0.24 sec\n",
      " - 分开你的一言 纪著第么我爱要的可爱女人 坏坏的让我疯狂的可爱女人 透坏的让我疯狂的可爱女人 透坏的让我疯\n",
      " - 不分开 在是处这样打我妈妈这可躲你的脑笑就问题默风 说不开的字写在西元前一深我一定你已你 不要走的太快就像\n",
      "epoch 150, perplexity 4.430330, time 0.24 sec\n",
      " - 分开 娘子的黑猫 谁的完美主义 太彻底 让我连恨都难以下笔 将真不抽我已无的可爱女人 坏坏的让我疯狂的可\n",
      " - 不分开 不要处这样打我妈妈 难道你手不会痛吗  上地的我不  让它会不过 几天都没有喝水也能活 脑袋瓜我一\n",
      "epoch 200, perplexity 2.285859, time 0.24 sec\n",
      " - 分开 什么我 谁是神枪阳 我心定拖之护 我该上好主动 我爱上好生活 静静不悄 你已经离开我 不知不觉 我\n",
      " - 不分开 不要再不样打我妈妈 我的伤口被 一定看着日落 一直了我 全你盯人防守 篮下禁区 快使用不截棍 哼哼\n",
      "epoch 250, perplexity 1.761581, time 0.26 sec\n",
      " - 分开 小养我 谁是神枪阳 巫师 他念念 有词的 呼酋长下诅咒 所我骷髅头 这也它 一诉我 印地安的传说 \n",
      " - 不分开 不要你 我跟眼打那力  我去着只了信命运 在谢翻心引力 野梭碰到你 漂亮 是去开义 戒指在哭泣 静\n"
     ]
    }
   ],
   "source": [
    "num_epochs, batch_size, lr, clipping_theta = 250, 32, 1e2, 1e-2\n",
    "pred_period, pred_len, prefixes = 50, 50, ['分开', '不分开']\n",
    "train_and_predict_rnn_gluon(model, num_hiddens, vocab_size, ctx,\n",
    "                            corpus_indices, idx_to_char, char_to_idx,\n",
    "                            num_epochs, num_steps, lr, clipping_theta, \n",
    "                            batch_size, pred_period, pred_len, prefixes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.6 (XPython)",
   "language": "python",
   "name": "xpython"
  },
  "language_info": {
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "version": "3.6.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
